import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.ensemble import RandomForestRegressor
from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error  # 添加 MAE
import matplotlib
import seaborn as sns
sns.set(style='whitegrid')
matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体显示中文
matplotlib.rcParams['axes.unicode_minus'] = False    # 正确显示负号


# ========== 1. 数据加载 ==========
df = pd.read_csv('../data/cleaned_data.csv', encoding='gbk')

# ========== 2. 基础特征 + 派生特征 ==========
features = [
    'weiliaoc', 'yaotouc', 'bilengjydS1', 'bilengjedS1',
    'bilengjsdI1', 'shuliaoKH', 'shuliaoIM', 'fengjizs', 'rehao'
]
target = 'shuliaoSM'

# 时间字段转换
df['shijian'] = pd.to_datetime(df['date'].astype(str) + ' ' + df['time'].astype(str), errors='coerce')

df['hour'] = df['shijian'].dt.hour
df['day'] = df['shijian'].dt.dayofweek
df['KH_IM_ratio'] = df['shuliaoKH'] / (df['shuliaoIM'] + 1e-6)

# 添加可用的新字段（如果存在）
for col in ['shuliaol', 'chumoslKH', 'chumoslSM']:
    if col in df.columns:
        features.append(col)

# 添加派生特征
features += ['hour', 'day', 'KH_IM_ratio']

# ========== 3. 数据准备 ==========
df = df[features + [target]].dropna()

# 编码类别型字段
for col in ['weiliaoc', 'yaotouc']:
    if df[col].dtype == 'object':
        df[col] = LabelEncoder().fit_transform(df[col])

X = df[features]
y = df[target]

# 划分训练测试集
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# ========== 4. 随机森林 + 参数调优 ==========
param_grid = {
    'n_estimators': [100, 200],
    'max_depth': [5, 10, 15],
    'min_samples_split': [2, 5],
    'min_samples_leaf': [1, 2]
}

grid_search = GridSearchCV(
    estimator=RandomForestRegressor(random_state=42),
    param_grid=param_grid,
    scoring='neg_root_mean_squared_error',
    cv=5,
    verbose=1,
    n_jobs=-1
)

grid_search.fit(X_train, y_train)
best_model = grid_search.best_estimator_

print("✅ 最佳参数:", grid_search.best_params_)

# ========== 5. 预测与评估 ==========
y_pred = best_model.predict(X_test)
rmse = np.sqrt(mean_squared_error(y_test, y_pred))
mae = mean_absolute_error(y_test, y_pred)
r2 = r2_score(y_test, y_pred)

print(f"RMSE（调参后）: {rmse:.4f}")
print(f"MAE: {mae:.4f}")  # ← 输出 MAE
print(f"R²（调参后）: {r2:.4f}")

# ✅ 新增预测输出
output_df = pd.DataFrame({
    'True_SM': y_test.values,
    'Predicted_SM': y_pred
})
output_df.to_csv('../data/prediction_results.csv', index=False, encoding='gbk')
print("📁 预测结果已保存至 prediction_results.csv")
# ========== 6. 特征重要性可视化 ==========
importances = best_model.feature_importances_
sorted_idx = np.argsort(importances)[::-1]

plt.figure(figsize=(10, 5))
plt.bar([features[i] for i in sorted_idx], importances[sorted_idx])
plt.title("Feature Importances (Random Forest)")
plt.xticks(rotation=45)
plt.tight_layout()
plt.grid(True)
plt.show()

# ========== 7. 预测效果对比图 ==========
plt.figure(figsize=(6, 6))
plt.scatter(y_test, y_pred, alpha=0.6, edgecolors='k')
plt.plot([y.min(), y.max()], [y.min(), y.max()], '--', color='gray')
plt.xlabel("True SM")
plt.ylabel("Predicted SM")
plt.title("预测值 vs 实际值")
plt.grid(True)
plt.tight_layout()
plt.show()

# ========== 8. 保存预测结果 ==========
output_df = pd.DataFrame({
    'True_SM': y_test.values,
    'Predicted_SM': y_pred
})
output_df['Error'] = output_df['Predicted_SM'] - output_df['True_SM']
output_df.to_csv('../data/prediction_results.csv', index=False, encoding='gbk')
print("📁 预测结果已保存至 prediction_results.csv")

# ========== 9. 可视化 ==========

# ① 预测值 vs 实际值散点图
plt.figure(figsize=(6,6))
plt.scatter(output_df['True_SM'], output_df['Predicted_SM'], alpha=0.6, edgecolors='k')
plt.plot([output_df['True_SM'].min(), output_df['True_SM'].max()],
         [output_df['True_SM'].min(), output_df['True_SM'].max()],
         '--', color='gray')
plt.xlabel("True SM")
plt.ylabel("Predicted SM")
plt.title("① 预测值 vs 实际值")
plt.grid(True)
plt.tight_layout()
plt.show()

# ② 预测误差分布图
plt.figure(figsize=(8,4))
sns.histplot(output_df['Error'], bins=30, kde=True, color='orange')
plt.axvline(0, color='gray', linestyle='--')
plt.title("② 预测误差分布（Predicted - True）")
plt.xlabel("误差")
plt.tight_layout()
plt.show()

# ③ 预测误差箱线图
plt.figure(figsize=(6,4))
sns.boxplot(x=output_df['Error'], color='skyblue')
plt.title("③ 预测误差箱线图")
plt.tight_layout()
plt.show()

# ④ 残差 vs 预测值图
plt.figure(figsize=(6,4))
plt.scatter(output_df['Predicted_SM'], output_df['Error'], alpha=0.5, edgecolors='k')
plt.axhline(0, linestyle='--', color='gray')
plt.xlabel("Predicted SM")
plt.ylabel("Residual (Pred - True)")
plt.title("④ 残差 vs 预测值")
plt.tight_layout()
plt.show()
